#!/usr/bin/env python3
"""
Simplified SfM Pipeline Script for Docker
This script runs the SfM pipeline with error handling and logging.
"""

import os
import sys
import argparse
from pathlib import Path

# Import hloc modules
from hloc import (
    extract_features,
    match_features,
    reconstruction,
    visualization,
    pairs_from_retrieval,
)


def run_sfm_pipeline(images_path, outputs_path=None):
    """Run the complete SfM pipeline."""
    
    # Setup paths
    images = Path(images_path)
    if outputs_path is None:
        outputs = Path("outputs/sfm/")
    else:
        outputs = Path(outputs_path)
    
    sfm_pairs = outputs / "pairs-netvlad.txt"
    sfm_dir = outputs / "sfm_superpoint+superglue"
    
    # Create output directory
    outputs.mkdir(parents=True, exist_ok=True)
    
    # Configuration
    retrieval_conf = extract_features.confs["netvlad"]
    feature_conf = extract_features.confs["superpoint_aachen"]
    matcher_conf = match_features.confs["superglue"]
    
    print("Starting SfM Pipeline...")
    print(f"Images: {images}")
    print(f"Outputs: {outputs}")
    
    # Check if images directory exists
    if not images.exists():
        print(f"Error: Images directory does not exist: {images}")
        return False
    
    try:
        # Step 1: Image retrieval
        print("Step 1: Extracting global descriptors...")
        retrieval_path = extract_features.main(retrieval_conf, images, outputs)
        pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=5)
        
        # Step 2: Local features
        print("Step 2: Extracting and matching local features...")
        feature_path = extract_features.main(feature_conf, images, outputs)
        match_path = match_features.main(
            matcher_conf, sfm_pairs, feature_conf["output"], outputs
        )
        
        # Step 3: Reconstruction
        print("Step 3: Running 3D reconstruction...")
        model = reconstruction.main(sfm_dir, images, sfm_pairs, feature_path, match_path)
        
        # Step 4: Visualization
        print("Step 4: Generating visualizations...")
        visualization.visualize_sfm_2d(model, images, color_by="visibility", n=5)
        visualization.visualize_sfm_2d(model, images, color_by="track_length", n=5)
        visualization.visualize_sfm_2d(model, images, color_by="depth", n=5)
        
        print("Pipeline completed successfully!")
        return True
        
    except Exception as e:
        print(f"Pipeline failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Main function with argument parsing."""
    parser = argparse.ArgumentParser(
        description="Run SfM pipeline on a set of images",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 run_sfm.py /path/to/images
  python3 run_sfm.py /path/to/images --output /path/to/output
  python3 run_sfm.py /media/sdisk/data/3dgs/20250731/South-Building/images
        """
    )
    
    parser.add_argument(
        "images_path",
        help="Path to the directory containing input images"
    )
    
    parser.add_argument(
        "--output", "-o",
        help="Path to the output directory (default: outputs/sfm/)",
        default=None
    )
    
    args = parser.parse_args()
    
    success = run_sfm_pipeline(args.images_path, args.output)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
